import asyncio
import json
import logging
from typing import Dict

import websockets, time


WS_URL = 'wss://ws.okx.com:8443/ws/v5/public'

logger = logging.root
if not logger.hasHandlers():
    logger.setLevel(logging.INFO)
    logger.addHandler(logging.StreamHandler())


async def main():
    while True:
        mark_price_details = {

        }
        index_price = None
        client = None
        try:
            async with websockets.connect(uri=WS_URL, ping_interval=5, timeout=10) as client:
                subscribes = {
                    'op': 'subscribe',
                    'args': [
                        {'channel': 'index-tickers', 'instId': 'BTC-USD'},
                        # {'channel': 'mark-price', 'instId': 'BTC-USD-SWAP'},
                        {'channel': 'mark-price', 'instId': 'BTC-USD-241227'},
                        {'channel': 'mark-price', 'instId': 'BTC-USD-250131'},
                        {'channel': 'mark-price', 'instId': 'BTC-USD-250328'},
                        {'channel': 'mark-price', 'instId': 'BTC-USD-250627'},
                    ]
                }
                await client.send(json.dumps(subscribes))
                resp = json.loads(await asyncio.wait_for(client.recv(), 5))
                if resp.get('event') == 'error':
                    raise ValueError(resp.get('msg'))
                logger.info(f'subscribe result is {resp}')

                while not client.closed:
                    try:
                        msg: Dict = json.loads(await client.recv())
                        channel = msg.get('arg', {}).get('channel')
                        data = msg.get('data',  [])
                        if not data:
                            continue
                        data = data[0]

                        if channel == 'index-tickers':
                            index_price = float(data['idxPx'])

                        elif channel == 'mark-price':
                            inst_id = msg.get('arg', {}).get('instId')
                            d = mark_price_details.setdefault(inst_id, {
                                'count': 0,
                                'arr': [0] * 10000,
                                'avg': 0,
                                'total': 0
                            })
                            if not index_price:
                                continue
                            p = float(data['markPx']) / index_price
                            abc = (p - d['avg']) / d['avg'] if d['avg'] else 0

                            if d['count'] >= 1000 and abs(abc) >= 0.002:
                                abc = round(p - d['avg'], 4)
                                print(f'gate: {abc}\t{time.time() * 1000 // 1}')

                            d['count'] += 1
                            idx: int = (d['count'] - 1) % len(d['arr'])
                            d['total'] = d['total'] - d['arr'][idx] + p
                            d['avg'] = d['total'] / min(d['count'], len(d['arr']))
                            d['arr'][idx] = p

                            if d['count'] % 50 == 0:
                                print(f't: {inst_id}\t{format(p, ".4f")}\t{format(d["avg"], ".4f")}\t{format(abc, "+.4f")}\t  {time.time() * 1000 // 1}')

                    except websockets.exceptions.ConnectionClosed:
                        print('okx ws connect is closed, will retry connect')
        except websockets.exceptions.ConnectionClosed:
            print('okx ws connect is closed, will retry connect')
        except Exception as e:
            logger.exception(f'error is {e}')
            # traceback.print_stack()
            # print(f'handle okx websocket failed: {traceback.format_stack()}')
        finally:
            if not client.closed:
                try:
                    await client.close()
                finally:
                    pass
            await asyncio.sleep(10)


if __name__ == '__main__':
    asyncio.run(main())

